from pocketflow import AsyncFlow
from typing import Dict, Any, Optional

from .microflow import(
    PrepareAction,
    StreamPlanAction,
    MCPAction,
    StreamAnswerAction,
    Done
)

from scorpio.core.common import get_logger
logger = get_logger(__name__)

default_max_iteration=10

# to-do: check mcp server and llm in initializing agent
class DefaultAgent:
    shared: Dict
    _async_flow: Optional[AsyncFlow] = None
    def __init__(self, shared: Dict, max_iteration: int=None):
        self.name="default"
        shared["context"]["max_iteration"]=max_iteration or default_max_iteration
        self.shared = shared
        self.prepare=PrepareAction()
        self.plan=StreamPlanAction()
        self.mcp = MCPAction()
        self.answer = StreamAnswerAction()
        self.done = Done()

    def create(self)->AsyncFlow:
        """
        Create and connect the nodes to form a complete agent flow.
        
        The flow works like this:
        1. DecideAction node decides whether to search or answer
        2. If search, go to SearchWeb node
        3. If answer, go to AnswerQuestion node
        4. After SearchWeb completes, go back to DecideAction
        
        Returns:
            Flow: A complete research agent flow
        """
        
        self.prepare - "failed" >> self.done
        self.prepare - "prepared" >> self.plan

        self.plan - "zero" >> self.done
        self.plan - "single" >> self.mcp
        self.mcp - "success" >> self.plan
        self.mcp - "failure" >> self.answer 

        # Create and return the flow, starting with the DecideAction node
        self._flow=AsyncFlow(start=self.prepare)
        return self._flow
    async def run(self):
        logger.info(f"start to run agent: {self.name}")
        return await self._flow.run_async(self.shared)
